package com.kyubi.common.model;

import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Map;
import java.util.Queue;
import java.util.Set;
import java.util.SortedMap;
import java.util.SortedSet;
import java.util.TreeMap;
import java.util.Vector;

import com.kyubi.common.grammar.Tree;
import com.kyubi.common.grammar.TreeNode;
import com.kyubi.tagger.taggers.StanfordTreeTagger;

public class Lattice {
	
	static final double UNKNOWN_PROB = 0.000001;
	
	public Lattice(){
		
	}
	public Lattice(Map<String, Double> channelProb, StanfordTreeTagger treeTagger){
		 m_ChannelProb = channelProb; 
		 m_Lattice = new HashMap<Integer , Map<String, LatticeNode> >();
		 m_Tagger = treeTagger;
	}
	public void buildLattice(String input, Map<String, Double> wordFreq, Map<String, Double> ruleFreq){
		int numLevels = input.split("[ ]+").length;
		Map<String , LatticeNode> levelMap = null;
		for(int k = 0; k < numLevels; ++k){
			levelMap= new HashMap<String, LatticeNode>();
			m_Lattice.put(k, levelMap);
		}
		//Vector<String> subString = getAllSubsequences(input);
		Vector<String> subString = new Vector<String>();
		subString.add(input);
		Vector<String> innerString = new Vector<String>();
		Tree lTree = new Tree();
		Tree rTree = new Tree();
		int cLevel = 0;
		while(true){
			if(subString.get(0).split("[ ]+").length == 2){
				break;
			}
			for(int i = 0; i < subString.size(); ++i){
				innerString.addAll(getAllSubsequences(subString.get(i), wordFreq, ruleFreq));
				lTree.parse(m_Tagger.parse(subString.get(i)));
	
				
				for(int j = 0; j < innerString.size(); ++j){
					TreeNode node = rTree.parse(m_Tagger.parse(innerString.get(j)));
					double prior = computePriorProb(innerString.get(j), node, wordFreq, ruleFreq);
					m_ChannelProb.put(innerString.get(j), prior);
					addToLattice(innerString.get(j), lTree.getID(), rTree.getID(), cLevel, prior);
				}
			}
			subString = (Vector<String>)innerString.clone();
			innerString.clear();
			cLevel++;
		}
		
		levelMap = m_Lattice.get(numLevels-2);
		Set<String> keys = levelMap.keySet();
		Iterator<String> iter = keys.iterator();
		while(iter.hasNext()){
			addToLattice("END", iter.next(), "END", numLevels-2, 0);
		}
		
		//special handling for the first node....
		levelMap = m_Lattice.get(0);
		TreeNode node = lTree.parse(m_Tagger.parse(input));
		LatticeNode lnode = levelMap.get(lTree.getID());
		double prior = computePriorProb(input, node, wordFreq, ruleFreq);
		lnode.m_NodeLabel = input;
		lnode.m_TreePrior = prior;
		m_ChannelProb.put(input, prior);
		levelMap.put(lTree.getID(), lnode);
		m_Lattice.put(0, levelMap);
		
		//should be removed...
		keys = m_ChannelProb.keySet();
		iter = keys.iterator();
		while(iter.hasNext()){
			String key = iter.next();
			System.out.println(key.split("[ ]+").length + " " + key +  " " + m_ChannelProb.get(key));
		}
		
	}
	public void doForward(){
		int level = m_Lattice.size();
		Map<String, LatticeNode> cLevelMap = null;
		Map<String, LatticeNode> nLevelMap = null;
		for(int k = 0; k < level-1; ++k){
			 cLevelMap = m_Lattice.get(k);
			 nLevelMap = m_Lattice.get(k+1);
			 Set<String> keys = nLevelMap.keySet();
			 Iterator<String> iter = keys.iterator();
			 while(iter.hasNext()){
				 String key = iter.next();
				 LatticeNode node = nLevelMap.get(key);
				 double sum = 0;
				 for(int i = 0; i < node.m_ParentNodes.size(); ++i){
					 String bigram = node.m_ParentNodes.get(i) + "_" + key;
					 LatticeNode pNode = cLevelMap.get(node.m_ParentNodes.get(i));
					 if(m_ChannelProb.containsKey(bigram)){
						 sum = sum + pNode.m_RAlpha * m_ChannelProb.get(bigram); 
					 }
					 else{
						 sum = sum + pNode.m_RAlpha * UNKNOWN_PROB;
					 }
				 }
				 node.m_LAlpha = sum;
				 node.m_RAlpha = node.m_LAlpha * node.m_TreePrior;
				 nLevelMap.put(key, node);				
			}
			m_Lattice.put(k+1, nLevelMap);
		}
	}
	public void doBackward(){
		int level = m_Lattice.size();
		Map<String, LatticeNode> cLevelMap = null;
		Map<String, LatticeNode> nLevelMap = null;
		for(int k = level-1; k > 0; --k){
			 cLevelMap = m_Lattice.get(k);
			 nLevelMap = m_Lattice.get(k-1);
			 Set<String> keys = nLevelMap.keySet();
			 Iterator<String> iter = keys.iterator();
			 while(iter.hasNext()){
				 String key = iter.next();
				 LatticeNode node = nLevelMap.get(key);
				 double sum = 0;
				 for(int i = 0; i < node.m_ChildNodes.size(); ++i){
					 String bigram = key + "_" + node.m_ChildNodes.get(i);
					 LatticeNode cNode = cLevelMap.get(node.m_ChildNodes.get(i));
					 if(m_ChannelProb.containsKey(bigram)){
						 sum = sum + cNode.m_LBeta * m_ChannelProb.get(bigram); 
					 }
					 else{
						 sum = sum + cNode.m_LBeta * UNKNOWN_PROB;
					 }					 
				 }
				 node.m_RBeta = sum;
				 node.m_LBeta = node.m_RBeta * node.m_TreePrior;
				 nLevelMap.put(key, node);								 
			}
			m_Lattice.put(k+1, nLevelMap);			 
		}
		
	}
	private void addToLattice(String input, String parent, String child, int level, double pPrior){
//		System.out.println(input + " " +pPrior);
		LatticeNode nodeP = null;
		LatticeNode nodeC = null;
		Map<String, LatticeNode> levelMap = null;
		
		levelMap = m_Lattice.get(level);
		if(levelMap.containsKey(parent)){
			nodeP = levelMap.get(parent);
		}
		else{
			nodeP = new LatticeNode();
		}
		nodeP.m_ChildNodes.add(child);
		levelMap.put(parent, nodeP);
		m_Lattice.put(level, levelMap);

		levelMap = m_Lattice.get(level+1);
		if(levelMap.containsKey(child)){
			nodeC = levelMap.get(child);
		}
		else{
			nodeC = new LatticeNode();
		}
		nodeC.m_ParentNodes.add(parent);
		nodeC.m_NodeLabel = input;
		nodeC.m_TreePrior = pPrior;
		levelMap.put(child, nodeC);
		m_Lattice.put(level+1, levelMap);		
	}
	public Vector<String> getAllSubsequences(String inputSentence, Map<String, Double> wordFreq, Map<String, Double> ruleFreq){
		Vector<String> outputSentences = new Vector<String>();
		String str = null;
		int fIndex = 0;
		int sIndex = inputSentence.indexOf(' ') + 1;
		
		while(sIndex < inputSentence.length()){
			str = inputSentence.substring(0, fIndex) + inputSentence.substring(sIndex, inputSentence.length());
			outputSentences.add(str);
			fIndex = sIndex;
			if(inputSentence.indexOf(' ', fIndex) != -1){
				sIndex = inputSentence.indexOf(' ', fIndex) + 1;
			}
			else{
				break;
			}
			
		}
		str = inputSentence.substring(0, fIndex);
		outputSentences.add(str);
		
		//this is a hack.....
		Tree tree = new Tree();		
		Map<Double, String> map = new TreeMap<Double, String>();
		for(int i = 0; i < outputSentences.size(); ++i){
			TreeNode node = tree.parse(m_Tagger.parse(outputSentences.get(i)));
			map.put(computePriorProb(outputSentences.get(i), node, wordFreq, ruleFreq), outputSentences.get(i));
		}

		Set<Double> keys = map.keySet();
		String value = null;
		for(Double key : keys) {
			value = map.get(key);
		}
		Vector<String> temp = new Vector<String>();
		temp.add(value);
		return temp;
		//return outputSentences;
	}
	private double computePriorProb(String input, TreeNode node, Map<String, Double> wordFreq, Map<String, Double> ruleFreq){
		double treePrior = 1;
		TreeNode outNode = null;
		String unigramRule = null;
		String bigramRule = new String();
		Queue<TreeNode> queue = new LinkedList<TreeNode>();
		queue.add(node);
		while(!queue.isEmpty()){
			outNode = queue.remove();
			if(outNode.getParent() != null){
				unigramRule = /*outNode.getParent().getHead() + "->" +*/ outNode.getHead(); 
			}
			else{
				unigramRule = outNode.getHead();
			}
			Vector<TreeNode> childNodes = outNode.getChildren();
			if(!childNodes.isEmpty()){
				for(int i = 0; i < childNodes.size(); ++i){
					bigramRule = bigramRule + childNodes.get(i).getHead();
					/*bigramRule = unigramRule + "_" + childNodes.get(i).getParent().getHead() + "->" + childNodes.get(i).getHead();
					if(ruleFreq.containsKey(bigramRule)){
						treePrior = treePrior * ruleFreq.get(bigramRule)/ruleFreq.get(unigramRule);
					}
					else{
						treePrior = treePrior * UNKNOWN_PROB;
					}*/
					queue.add(childNodes.get(i));
				}
				bigramRule = unigramRule + "_" + unigramRule + "->" + bigramRule;
				if(ruleFreq.containsKey(bigramRule)){
					treePrior = treePrior * ruleFreq.get(bigramRule)/ruleFreq.get(unigramRule);
				}
				else{
					treePrior = treePrior * UNKNOWN_PROB;
				}
				bigramRule = "";
			}
		}
		String[] words = input.split("[ ]+");
		for(int i = 0; i < words.length-1; ++i){
			String bigram = words[i] + "_" + words[i+1];
			if(wordFreq.containsKey(bigram)){
				treePrior = treePrior * wordFreq.get(bigram)/wordFreq.get(words[i]);
			}
			else{
				treePrior = treePrior * UNKNOWN_PROB;
			}
			
		}
		return treePrior;
	}
	private Map<Integer , Map<String, LatticeNode> > m_Lattice;
	private StanfordTreeTagger m_Tagger;
	private Map<String, Double> m_ChannelProb;
}
